#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Visualize the mesh structure to show the wing surface and fluid domain.
This helps understand why the mesh looks like a "cavity".
"""

import pyvista as pv
import numpy as np
import os
import sys

# Add the dragonfly_simulation directory to Python path
DRAGONFLY_SIM_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(DRAGONFLY_SIM_ROOT)

import config

def visualize_mesh_structure():
    """Create a detailed visualization of the actual mesh structure."""
    
    mesh_file = os.path.join(DRAGONFLY_SIM_ROOT, config.MESH_FILE)
    
    if not os.path.exists(mesh_file):
        print(f"ERROR: Mesh file not found at {mesh_file}")
        print("Please run 'python generate_mesh.py' first!")
        return
    
    print("Loading mesh...")
    mesh = pv.read(mesh_file)
    print(f"Mesh loaded: {mesh.n_points} points, {mesh.n_cells} cells")
    
    # Create plotter (off_screen=True for saving images)
    plotter = pv.Plotter(window_size=[2400, 1400], off_screen=True)
    plotter.add_text("Dragonfly CFD Mesh Structure - Actual Mesh", position='upper_left', font_size=20)
    
    # Show the actual mesh structure with edges visible
    print("Analyzing mesh structure...")
    
    # Extract surface to show boundaries
    surface = mesh.extract_surface()
    print(f"Surface extracted: {surface.n_cells} boundary faces")
    
    # Show the mesh surface with visible edges
    plotter.add_mesh(surface, color='lightblue', opacity=0.6, 
                    show_edges=True, edge_color='navy', line_width=0.5,
                    label='Mesh Boundaries')
    
    # Add multiple slices to show internal structure
    print("Creating cross-sectional slices...")
    
    # Y=0 slice (side view)
    slice_y = mesh.slice(normal='y', origin=(0, 0, 0))
    plotter.add_mesh(slice_y, show_edges=True, opacity=0.8, 
                    edge_color='darkred', line_width=1.0, color='lightcoral',
                    label='Y=0 Slice (Side View)')
    
    # Z=0 slice (top view)
    slice_z = mesh.slice(normal='z', origin=(0, 0, 0))
    plotter.add_mesh(slice_z, show_edges=True, opacity=0.7, 
                    edge_color='darkgreen', line_width=1.0, color='lightgreen',
                    label='Z=0 Slice (Top View)')
    
    # Find and highlight potential wing region
    center = surface.cell_centers()
    distances = np.linalg.norm(center.points, axis=1)
    wing_threshold = 0.05  # 5cm from origin
    
    wing_cells = np.where(distances < wing_threshold)[0]
    print(f"Potential wing cells found: {len(wing_cells)}")
    
    if len(wing_cells) > 0:
        # Highlight wing region
        selection = np.zeros(surface.n_cells, dtype=bool)
        selection[wing_cells] = True
        wing_surface = surface.extract_cells(selection)
        
        plotter.add_mesh(wing_surface, color='red', opacity=1.0, 
                        show_edges=True, edge_color='darkred', line_width=2.0,
                        label='Wing Region')
        print(f"Wing region highlighted: {wing_surface.n_cells} cells")
    else:
        print("No wing cells found near origin - showing mesh bounds")
        # Show mesh bounds
        bounds = mesh.bounds
        print(f"Mesh bounds: X[{bounds[0]:.3f}, {bounds[1]:.3f}], Y[{bounds[2]:.3f}, {bounds[3]:.3f}], Z[{bounds[4]:.3f}, {bounds[5]:.3f}]")
    
    # Add coordinate axes
    plotter.add_axes()
    
    # Add annotations with mesh statistics
    mesh_info = f"Mesh Statistics:\n"
    mesh_info += f"Points: {mesh.n_points:,}\n"
    mesh_info += f"Cells: {mesh.n_cells:,}\n"
    mesh_info += f"Boundary Faces: {surface.n_cells:,}\n"
    mesh_info += f"Cell Types: {len(np.unique(mesh.celltypes))}"
    
    plotter.add_text(mesh_info, position='lower_left', font_size=12)
    
    # Add legend
    plotter.add_legend()
    
    # Set camera for good view
    plotter.camera_position = 'xy'
    plotter.camera.zoom(2.0)  # Zoom in to see mesh details better
    
    # Save the plot as PNG
    output_path = "mesh_structure_visualization.png"
    plotter.screenshot(output_path)
    print(f"\nMesh structure visualization saved to: {output_path}")
    plotter.close()

def create_raw_mesh_visualization():
    """Create a visualization showing the raw mesh structure clearly."""
    
    mesh_file = os.path.join(DRAGONFLY_SIM_ROOT, config.MESH_FILE)
    
    if not os.path.exists(mesh_file):
        print(f"ERROR: Mesh file not found at {mesh_file}")
        return
    
    print("Loading mesh for raw visualization...")
    mesh = pv.read(mesh_file)
    print(f"Mesh loaded: {mesh.n_points} points, {mesh.n_cells} cells")
    
    # Create plotter for raw mesh
    plotter = pv.Plotter(window_size=[2400, 1400], off_screen=True)
    plotter.add_text("Raw Mesh Structure (All Elements)", position='upper_left', font_size=20)
    
    # Show a subset of the mesh to make it visible
    # Take every nth cell to avoid overcrowding
    n_cells = mesh.n_cells
    step = max(1, n_cells // 5000)  # Show max 5000 cells
    cell_indices = list(range(0, n_cells, step))
    
    cell_subset = mesh.extract_cells(cell_indices)
    print(f"Showing {cell_subset.n_cells} cells out of {n_cells} total")
    
    # Add the raw mesh with edges
    plotter.add_mesh(cell_subset, show_edges=True, opacity=0.3, 
                    color='lightgray', edge_color='black', line_width=0.5,
                    label='Raw Mesh Elements')
    
    # Add coordinate axes
    plotter.add_axes()
    
    # Add mesh info
    plotter.add_text(f"Showing {cell_subset.n_cells:,} of {n_cells:,} total cells\n" +
                    f"Mesh type: Tetrahedral elements", 
                    position='lower_left', font_size=12)
    
    # Set camera
    plotter.camera_position = 'xy'
    plotter.camera.zoom(2.0)  # Zoom in to see mesh details better
    
    # Save the plot as PNG
    output_path = "raw_mesh_structure.png"
    plotter.screenshot(output_path)
    print(f"\nRaw mesh visualization saved to: {output_path}")
    plotter.close()

def create_conceptual_diagram():
    """Create a simple 2D conceptual diagram."""
    
    plotter = pv.Plotter(window_size=[1200, 800], off_screen=True)
    plotter.add_text("CFD Mesh Concept: Fluid Domain Around Wing", 
                    position='upper_left', font_size=18)
    
    # Create a 2D slice representation
    # Outer box (fluid domain)
    outer_box = pv.Box(bounds=(-0.4, 0.4, -0.35, 0.35, -0.01, 0.01))
    
    # Wing representation (ellipse-like)
    theta = np.linspace(0, np.pi/2, 50)
    x = 0.035 * np.cos(theta)
    y = 0.02 * np.sin(theta)
    z = np.zeros_like(x)
    
    # Close the wing shape
    x = np.concatenate([x, [0], [x[0]]])
    y = np.concatenate([y, [0], [y[0]]])
    z = np.concatenate([z, [0], [z[0]]])
    
    points = np.column_stack((x, y, z))
    # Create lines connecting consecutive points
    n_points = len(points)
    lines = []
    for i in range(n_points - 1):
        lines.extend([2, i, i + 1])  # 2 points per line, then the point indices
    
    wing = pv.PolyData(points, lines=lines)
    
    # Add to plotter
    plotter.add_mesh(outer_box, color='lightblue', opacity=0.3, 
                    label='Fluid Domain (Air)')
    plotter.add_mesh(wing, color='red', line_width=5, 
                    label='Wing Boundary')
    
    # Add flow arrows
    arrow_start = pv.PolyData([[-0.35, 0, 0]])
    arrow = arrow_start.glyph(orient=[1, 0, 0], scale=False, 
                             factor=0.1, geom=pv.Arrow())
    plotter.add_mesh(arrow, color='green', label='Flow Direction')
    
    # Add text annotations
    plotter.add_text("We mesh THIS region\n(the air)", 
                    position=(0.55, 0.7), font_size=14)
    plotter.add_text("NOT the wing itself", 
                    position=(0.45, 0.3), font_size=14)
    
    # Add legend
    plotter.add_legend()
    
    # Set view
    plotter.view_xy()
    plotter.camera.zoom(1.3)
    
    # Save image
    output_path = "mesh_concept_diagram.png"
    plotter.screenshot(output_path)
    print(f"\nConceptual diagram saved to: {output_path}")
    plotter.close()

if __name__ == "__main__":
    print("=== Mesh Structure Visualization ===")
    print("\n1. Creating conceptual diagram...")
    create_conceptual_diagram()
    
    print("\n2. Creating raw mesh visualization...")
    create_raw_mesh_visualization()
    
    print("\n3. Creating detailed mesh structure visualization...")
    visualize_mesh_structure() 